import os
import os.path as osp
import shutil
import threading
import cv2
from PyQt5.QtCore import Qt
from PyQt5.QtGui import QIcon, QPixmap, QFont, QFontDatabase
from PyQt5.QtWidgets import (QTabWidget, QWidget, QVBoxLayout, QHBoxLayout, 
                            QLabel, QPushButton, QFileDialog, QMessageBox,
                            QTableWidget, QTableWidgetItem, QHeaderView)

from utils.detector import Detector
from torch.serialization import add_safe_globals
from torch.serialization import safe_globals
from ultralytics.nn.tasks import DetectionModel
from ultralytics import YOLO

from ultralytics import YOLO
import cv2
import numpy as np

def detect_video(self, source, output_callback, stop_event, output_size=480, conf_thres=0.25, iou_thres=0.45, max_det=1000):
    """基于Ultralytics YOLOv8的视频检测

    Args:
        source: 视频源（文件路径或摄像头索引字符串）
        output_callback: 输出回调函数，用于实时显示结果 (frame, detections)
        stop_event: 线程停止事件
        output_size: 输出图像尺寸
        conf_thres: 置信度阈值
        iou_thres: IOU阈值
        max_det: 最大检测数量
    """
    model = self.model  # 假设模型已加载为YOLO对象，self.model = YOLO(...)
    cap = cv2.VideoCapture(int(source) if source.isnumeric() else source)

    while cap.isOpened():
        ret, frame = cap.read()
        if not ret or (stop_event and stop_event.is_set()):
            break

        # 推理
        results = model.predict(frame,
                                conf=conf_thres,
                                iou=iou_thres,
                                max_det=max_det,
                                device=self.device if hasattr(self, 'device') else 'cpu',
                                verbose=False)

        result = results[0]  # 单帧结果

        # 画框
        annotated_frame = result.plot()

        # 组织检测结果字典，符合原来格式
        detections = {
            'type': [result.names[int(cls)] for cls in result.boxes.cls.cpu().numpy()] if result.boxes else [],
            'box': result.boxes.xyxy.cpu().numpy().tolist() if result.boxes else [],
            'confidence': result.boxes.conf.cpu().numpy().tolist() if result.boxes else []
        }

        # 调整输出尺寸
        resize_scale = output_size / annotated_frame.shape[0]
        annotated_frame = cv2.resize(annotated_frame, (0, 0), fx=resize_scale, fy=resize_scale)

        # 回调输出
        if output_callback:
            output_callback(annotated_frame, detections)

        # cv2.waitKey(1) 处理键盘事件（非必须，根据你的程序决定）
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break

    cap.release()

def parse_yolo_detections(results):
    """将 YOLOv8 的结果转换为标准 detections 字典"""
    detections = {
        "type": [],
        "box": [],
        "confidence": []
    }

    if not results:
        return detections

    result = results[0]  # 单张图像

    boxes = result.boxes
    if boxes is None:
        return detections

    for cls_id, box, conf in zip(boxes.cls, boxes.xyxy, boxes.conf):
        detections["type"].append(result.names[int(cls_id)])
        detections["box"].append([int(v) for v in box.cpu().numpy()])
        detections["confidence"].append(float(conf.cpu().numpy()))
    
    return detections


class MainWindow(QTabWidget):
    """主窗口类"""

    def __init__(self):
        """初始化主窗口"""
        super().__init__()
        
        # 设置窗口基本属性
        self.setWindowTitle('井盖隐患智能识别')
        self.resize(1200, 800)
        self.setWindowIcon(QIcon("images/UI/lufei.png"))
        
        # 检测相关参数
        self.output_size = 480
        self.img2predict = ""
        self.device = 'cpu'
        self.vid_source = '0'  # 初始设置为摄像头
        self.webcam = True
        
        # 初始化停止事件
        self.stopEvent = threading.Event()
        self.stopEvent.clear()
        
        # 创建临时目录
        os.makedirs("images/tmp", exist_ok=True)
        
        # 加载模型
        with safe_globals([DetectionModel]):
            self.model = YOLO("runs/train/exp/weights/best.pt")
            # self.detector = Detector(weights="runs/train/exp/weights/best.pt", device=self.device)
            print(self.model.names)  # 查看模型的类别名

        
        
        # 初始化UI
        self.initUI()
        
        # 重置视频界面
        self.reset_vid()
        
    def initUI(self):
        """初始化用户界面"""
        # 设置字体
        font_title = self.createFont('楷体', 16)
        font_main = self.createFont('楷体', 14)
        
        # 创建图片检测界面
        img_detection_widget = self.createImageDetectionUI(font_title, font_main)
        
        # 创建视频检测界面
        vid_detection_widget = self.createVideoDetectionUI(font_title, font_main)
        
        # 添加选项卡
        self.addTab(img_detection_widget, '图片检测')
        self.addTab(vid_detection_widget, '视频检测')
        
    def createFont(self, family, size):
        """创建字体

        Args:
            family: 字体名称
            size: 字体大小

        Returns:
            字体对象
        """
        from PyQt5.QtGui import QFont
        # 使用更通用的字体
        common_font = 'Microsoft YaHei' if 'Microsoft YaHei' in [f for f in QFontDatabase().families()] else 'Arial'
        font = QFont(common_font, size)
        return font
    
    def createImageDetectionUI(self, font_title, font_main):
        """创建图片检测UI

        Args:
            font_title: 标题字体
            font_main: 主要字体

        Returns:
            图片检测界面控件
        """
        # 创建图片检测控件
        img_detection_widget = QWidget()
        img_detection_layout = QVBoxLayout()
        
        # 创建标题
        img_detection_title = QLabel("图片识别功能")
        img_detection_title.setFont(font_title)
        img_detection_title.setAlignment(Qt.AlignCenter)
        
        # 创建图片展示区域
        mid_img_widget = QWidget()
        mid_img_layout = QHBoxLayout()
        
        self.left_img = QLabel()
        self.right_img = QLabel()
        self.left_img.setPixmap(QPixmap("images/UI/up.jpeg"))
        self.right_img.setPixmap(QPixmap("images/UI/right.jpeg"))
        self.left_img.setAlignment(Qt.AlignCenter)
        self.right_img.setAlignment(Qt.AlignCenter)
        
        mid_img_layout.addWidget(self.left_img)
        mid_img_layout.addStretch(0)
        mid_img_layout.addWidget(self.right_img)
        mid_img_widget.setLayout(mid_img_layout)
        
        # 创建检测结果表格
        result_table_widget = QWidget()
        result_table_layout = QVBoxLayout()
        
        result_table_title = QLabel("检测结果")
        result_table_title.setFont(font_title)
        result_table_title.setAlignment(Qt.AlignCenter)
        
        self.result_table = QTableWidget(0, 3)
        self.result_table.setHorizontalHeaderLabels(["隐患类型", "区域", "置信度"])
        self.result_table.horizontalHeader().setSectionResizeMode(QHeaderView.Stretch)
        self.result_table.setFont(font_main)
        
        result_table_layout.addWidget(result_table_title)
        result_table_layout.addWidget(self.result_table)
        result_table_widget.setLayout(result_table_layout)
        
        # 创建按钮
        up_img_button = self.createStyledButton("上传图片", font_main)
        det_img_button = self.createStyledButton("开始检测", font_main)
        
        # 设置按钮样式
        up_img_button.setStyleSheet(
            "QPushButton {"
            "    background-color: #2196F3;"
            "    color: white;"
            "    border-radius: 8px;"
            "    padding: 10px 20px;"
            "    min-width: 120px;"
            "}"
            "QPushButton:hover {"
            "    background-color: #1976D2;"
            "}"
        )
        
        det_img_button.setStyleSheet(
            "QPushButton {"
            "    background-color: #4CAF50;"
            "    color: white;"
            "    border-radius: 8px;"
            "    padding: 10px 20px;"
            "    min-width: 120px;"
            "}"
            "QPushButton:hover {"
            "    background-color: #388E3C;"
            "}"
        )
        
        # 创建按钮布局
        button_layout = QHBoxLayout()
        button_layout.addWidget(up_img_button)
        button_layout.addWidget(det_img_button)
        button_layout.setSpacing(30)  # 增加按钮间距
        button_layout.setContentsMargins(0, 5, 0, 20)  # 减少顶部边距，保持底部边距
        
        # 连接按钮信号
        up_img_button.clicked.connect(self.upload_img)
        det_img_button.clicked.connect(self.detect_img)
        
        # 将按钮组加入主布局
        img_detection_layout.addWidget(img_detection_title, alignment=Qt.AlignCenter)
        img_detection_layout.addWidget(mid_img_widget, alignment=Qt.AlignCenter)
        img_detection_layout.addWidget(result_table_widget)
        
        # 创建一个新的水平布局容器
        button_container = QWidget()
        button_container.setLayout(button_layout)
        
        # 创建一个垂直布局来包含按钮容器并设置对齐方式
        vertical_container = QWidget()
        vertical_layout = QVBoxLayout()
        vertical_layout.addWidget(button_container, alignment=Qt.AlignTop | Qt.AlignHCenter)
        vertical_container.setLayout(vertical_layout)
        
        img_detection_layout.addWidget(vertical_container)
        img_detection_layout.setSpacing(25)
        
        # 强制更新布局
        img_detection_layout.activate()
        img_detection_widget.setLayout(img_detection_layout)
        
        # 添加调试日志
        import logging
        logging.basicConfig(level=logging.DEBUG)
        logging.debug("按钮布局参数: %s", button_layout.getContentsMargins())
        logging.debug("主布局间距: %s", img_detection_layout.spacing())
        
        return img_detection_widget
    
    def createVideoDetectionUI(self, font_title, font_main):
        """创建视频检测UI

        Args:
            font_title: 标题字体
            font_main: 主要字体

        Returns:
            视频检测界面控件
        """
        # 创建视频检测控件
        vid_detection_widget = QWidget()
        vid_detection_layout = QVBoxLayout()
        
        # 创建标题
        vid_title = QLabel("视频检测功能")
        vid_title.setFont(font_title)
        vid_title.setAlignment(Qt.AlignCenter)
        
        # 创建视频展示区域
        vid_img_container = QWidget()
        vid_img_layout = QVBoxLayout()
        
        self.vid_img = QLabel()
        self.vid_img.setPixmap(QPixmap("images/UI/up.jpeg"))
        self.vid_img.setAlignment(Qt.AlignCenter)
        
        vid_img_layout.addWidget(self.vid_img)
        vid_img_container.setLayout(vid_img_layout)
        
        # 创建检测结果表格
        vid_result_table_widget = QWidget()
        vid_result_table_layout = QVBoxLayout()
        
        vid_result_table_title = QLabel("检测结果")
        vid_result_table_title.setFont(font_title)
        vid_result_table_title.setAlignment(Qt.AlignCenter)
        
        self.vid_result_table = QTableWidget(0, 3)
        self.vid_result_table.setHorizontalHeaderLabels(["隐患类型", "区域", "置信度"])
        self.vid_result_table.horizontalHeader().setSectionResizeMode(QHeaderView.Stretch)
        self.vid_result_table.setFont(font_main)
        
        vid_result_table_layout.addWidget(vid_result_table_title)
        vid_result_table_layout.addWidget(self.vid_result_table)
        vid_result_table_widget.setLayout(vid_result_table_layout)
        
        # 创建按钮
        self.webcam_detection_btn = self.createStyledButton("摄像头实时监测", font_main)
        self.mp4_detection_btn = self.createStyledButton("视频文件检测", font_main)
        self.vid_stop_btn = self.createStyledButton("停止检测", font_main)

        # 设置按钮样式
        for btn in [self.webcam_detection_btn, self.mp4_detection_btn, self.vid_stop_btn]:
            btn.setStyleSheet(
                "QPushButton {"
                "    background-color: #2196F3;"
                "    color: white;"
                "    border-radius: 8px;"
                "    padding: 10px 20px;"
                "    min-width: 120px;"
                "}"
                "QPushButton:hover {"
                "    background-color: #1976D2;"
                "}"
            )

        # 连接按钮信号
        self.webcam_detection_btn.clicked.connect(self.open_cam)
        self.mp4_detection_btn.clicked.connect(self.open_mp4)
        self.vid_stop_btn.clicked.connect(self.close_vid)

        # 创建按钮布局
        vid_button_layout = QHBoxLayout()
        vid_button_layout.addWidget(self.webcam_detection_btn)
        vid_button_layout.addWidget(self.mp4_detection_btn)
        vid_button_layout.addWidget(self.vid_stop_btn)
        vid_button_layout.setSpacing(30)  # 按钮间距
        vid_button_layout.setContentsMargins(0, 5, 0, 20)  # 边距设置

        # 创建垂直容器来包含按钮布局
        vid_buttons_container = QWidget()
        vid_buttons_container.setLayout(vid_button_layout)

        # 组装布局
        vid_detection_layout.addWidget(vid_title)
        vid_detection_layout.addWidget(vid_img_container)
        vid_detection_layout.addWidget(vid_result_table_widget)
        vid_detection_layout.addWidget(vid_buttons_container)
        vid_detection_widget.setLayout(vid_detection_layout)
        
        return vid_detection_widget
    
    def createStyledButton(self, text, font):
        """创建样式化按钮

        Args:
            text: 按钮文本
            font: 按钮字体

        Returns:
            样式化的按钮对象
        """
        button = QPushButton(text)
        button.setFont(font)
        button.setStyleSheet("""
            QPushButton{color:white}
            QPushButton:hover{background-color: rgb(2,110,180);}
            QPushButton{background-color:rgb(48,124,208)}
            QPushButton{border:2px}
            QPushButton{border-radius:5px}
            QPushButton{padding:5px 5px}
            QPushButton{margin:5px 5px}
        """)
        return button
    
    
    
    def update_detection_table(self, table, detections):
        """更新检测结果表格

        Args:
            table: 表格控件
            detections: 检测结果
        """
        # 清空表格
        table.setRowCount(0)
        
        # 填充检测结果
        if detections and len(detections['type']) > 0:
            for i in range(len(detections['type'])):
                row_position = table.rowCount()
                table.insertRow(row_position)
                
                # 隐患类型
                type_item = QTableWidgetItem(detections['type'][i])
                table.setItem(row_position, 0, type_item)
                
                # 区域 (x1, y1, x2, y2)
                box = detections['box'][i]
                box_text = f"({box[0]},{box[1]}) - ({box[2]},{box[3]})"
                box_item = QTableWidgetItem(box_text)
                table.setItem(row_position, 1, box_item)
                
                # 置信度
                conf = float(detections['confidence'][i])
                conf_text = f"{conf:.2f}"
                conf_item = QTableWidgetItem(conf_text)
                table.setItem(row_position, 2, conf_item)
    
    def upload_img(self):
        """上传图片"""
        # 打开文件对话框选择图片
        file_path, _ = QFileDialog.getOpenFileName(
            self, '选择图片', '', 'Image Files(*.jpg *.png)'
        )
        
        if file_path:
            # 保存图片路径
            self.img2predict = file_path
            
            # 显示原始图片
            pixmap = QPixmap(file_path)
            self.left_img.setPixmap(pixmap.scaled(
                self.left_img.size(),
                Qt.KeepAspectRatio,
                Qt.SmoothTransformation
            ))
            
            # 清空结果图片
            self.right_img.setPixmap(QPixmap("images/UI/right.jpeg"))
            
            # 清空检测结果表格
            self.result_table.setRowCount(0)

    def detect_img(self):
        """检测图片"""
        if not self.img2predict:
            QMessageBox.warning(self, "请上传", "请先上传图片再进行检测")
            return
            
        results = self.model.predict(source=self.img2predict, conf=0.25, iou=0.45, device='cpu' )
        detections = parse_yolo_detections(results)

        for r in results:
            print("预测类别编号（cls）：", r.boxes.cls.cpu().numpy())
        result_image = results[0].plot()
        # 保存和显示结果
        cv2.imwrite("images/tmp/single_result.jpg", result_image)
        pixmap = QPixmap("images/tmp/single_result.jpg")
        self.right_img.setPixmap(pixmap.scaled(
            self.right_img.size(),
            Qt.KeepAspectRatio,
            Qt.SmoothTransformation
        ))
        
        # 更新检测结果表格
        self.update_detection_table(self.result_table, detections)
    
    def open_cam(self):
        """打开摄像头"""
        # 禁用按钮防止重复点击
        self.webcam_detection_btn.setEnabled(False)
        self.mp4_detection_btn.setEnabled(False)
        self.vid_stop_btn.setEnabled(True)
        
        # 设置视频源为摄像头
        self.vid_source = '0'
        self.webcam = True
        
        # 清空表格
        self.vid_result_table.setRowCount(0)
        
        # 启动检测线程
        th = threading.Thread(target=self.detect_vid)
        th.start()
    
    def open_mp4(self):
        """打开视频文件"""
        # 选择视频文件
        fileName, fileType = QFileDialog.getOpenFileName(self, '选择视频', '', '*.mp4 *.avi')
        if fileName:
            # 禁用按钮防止重复点击
            self.webcam_detection_btn.setEnabled(False)
            self.mp4_detection_btn.setEnabled(False)
            self.vid_stop_btn.setEnabled(True)
            
            # 设置视频源为文件
            self.vid_source = fileName
            self.webcam = False
            
            # 清空表格
            self.vid_result_table.setRowCount(0)
            
            # 启动检测线程
            th = threading.Thread(target=self.detect_vid)
            th.start()
    
    def detect_vid(self):
        """视频检测主函数"""
        # 定义回调函数用于更新UI
        def update_ui(frame, detections):
            cv2.imwrite("images/tmp/single_result_vid.jpg", frame)
            self.vid_img.setPixmap(QPixmap("images/tmp/single_result_vid.jpg"))
            
            # 更新检测结果表格
            self.update_detection_table(self.vid_result_table, detections)
        
        # 启动视频检测
        detect_video(
            self=self,
    source=self.vid_source,
    output_callback=update_ui,
    stop_event=self.stopEvent,
    output_size=self.output_size,
    conf_thres=0.25,
    iou_thres=0.45,
    max_det=1000
)
        
        # 检测结束后重置UI
        self.reset_vid()
    
    def reset_vid(self):
        """重置视频界面"""
        self.webcam_detection_btn.setEnabled(True)
        self.mp4_detection_btn.setEnabled(True)
        self.vid_img.setPixmap(QPixmap("images/UI/up.jpeg"))
        self.vid_source = '0'
        self.webcam = True
    
    def close_vid(self):
        """关闭视频检测"""
        self.stopEvent.set()
        self.reset_vid()
    
    def closeEvent(self, event):
        """窗口关闭事件处理"""
        reply = QMessageBox.question(
            self,
            '退出',
            "确定要退出程序吗？",
            QMessageBox.Yes | QMessageBox.No,
            QMessageBox.No
        )
        
        if reply == QMessageBox.Yes:
            self.close()
            event.accept()
        else:
            event.ignore() 
            
